import os
import PIL
import torch
import numpy as np
# import matplotlib.pyplot as plt

from os import listdir
from os.path import join
from PIL import Image
from utils.transform import itensity_normalize
from torch.utils.data.dataset import Dataset


class coronary_dataset(Dataset):
    def __init__(self, dataset_folder='./data', train_type='train', transform=None):
        self.transform = transform
        self.train_type = train_type

        # with open(join(dataset_folder, 'coronaryDataset.txt'), 'r') as f:
        #     self.image_all_list = f.readlines()
        # one_tenth = len(self.image_all_list) // 10
        # if self.train_type == 'train':
        #     self.img_list = self.image_all_list[0: one_tenth * 7]
        # elif self.train_type == 'validation':
        #     self.img_list = self.image_all_list[one_tenth * 7: one_tenth * 8]
        # elif self.train_type == 'test':
        #     self.img_list = self.image_all_list[one_tenth * 8: len(self.image_all_list)]

        if self.train_type == "train":
            with open("./data/train_filtered.txt", 'r') as f:
               self.img_list = f.readlines()
        elif self.train_type == "val":
            with open('./data/valid_filtered.txt', 'r') as f:
                self.img_list = f.readlines()
        elif self.train_type == "test":
            with open('./data/test_filtered.txt', 'r') as f:
                self.img_list = f.readlines()

    def __getitem__(self, item: int):
        image_name = self.img_list[item].replace('\n', '')
        image = Image.open('/media/handewei/数据集/medical/coronary_seg_data/mix/src/' + image_name)
        label = Image.open('/media/handewei/数据集/medical/coronary_seg_data/mix/label/' + image_name)

        sample = {'image': image, 'label': label}

        if self.transform is not None:
            # TODO: transformation to argument datasets
            sample = self.transform(sample, self.train_type)

        return sample['image'], sample['label'], self.img_list[item]

    def __len__(self):
        return len(self.img_list)

